# coding=utf-8

"""
author = jamon
"""

import asyncio
import random
from asyncio import Future
from typing import Dict, List, Union

from obespoir.base.common_define import ConnectionStatus, RpcMsgType
from obespoir.base.global_object import GlobalObject
from obespoir.base.ob_exception import ObException
from obespoir.config import NodeType
from obespoir.share.espoir_log import EspoirLog
from obespoir.share.singleton import Singleton


class RpcConnectionManager(EspoirLog, metaclass=Singleton):

    def __init__(self):
        # {conn_name:
        # {status: int(0:未连接，1：连接中，2：连接断开), conn: connect object
        # }}
        super().__init__()
        self.conns = {}

        # 存放节点类型信息 {int: [conn_name1, conn_name2, ...]}
        self.type_dict: Dict[str, List] = {}
        self.future_dict: Dict[str, Future] = {}

    @staticmethod
    def gen_node_name(host, port):
        return GlobalObject.gen_id(host, port)

    def add_future(self, seq: str, future: Future):
        self.future_dict[seq] = future

    def set_future_result_and_close(self, seq, result):
        self.future_dict[seq].set_result(result)
        del self.future_dict[seq]

    def get_connection(self, host, port):
        name = RpcConnectionManager.gen_node_name(host, port)
        if name not in self.conns.keys() or \
                ConnectionStatus.ESTABLISHED != self.conns[name]["status"]:
            return None
        return self.conns[name]["conn"]

    def get_available_connection(self, node_type: Union[str, NodeType]):
        """
        返回一个可用的连接名
        :param node_type:
        :return: node name
        """
        if isinstance(node_type, NodeType):
            node_type = node_type.value
        cur_nodes = self.type_dict.get(node_type, [])
        available_nodes = []
        self.debug("start", node_type=node_type, type_dict=self.type_dict,
                   all_conns=self.conns, cur_nodes=cur_nodes)
        for name in cur_nodes:
            if ConnectionStatus.ESTABLISHED == self.conns[name]["status"]:
                available_nodes.append(name)
        return random.choice(available_nodes) if available_nodes else None

    def store_connection(self, host, port, connect,
                         status=ConnectionStatus.LOSE):
        self.debug("start", host=host, port=port, connect=connect, status=status)
        self.conns[RpcConnectionManager.gen_node_name(host, port)] = {
            "status": status, "conn": connect, "host": host, "port": port}

    def lost_connection(self, host, port):
        name = RpcConnectionManager.gen_node_name(host, port)
        if name not in self.conns.keys():
            self.warn("can not find connection", name=name)
            return 0
        self.conns[name]["conn"].transport.close()
        self.conns[name]["status"] = ConnectionStatus.LOSE

    def remove_connection(self, conn_name):
        if conn_name not in self.conns.keys():
            return 0
        self.conns[conn_name]["conn"].transport.close()
        self.conns.pop(conn_name)
        return 1

    def add_type_node(self, node_type: str, host: str, port: int):
        if isinstance(node_type, NodeType):
            node_type = node_type.value
        name = RpcConnectionManager.gen_node_name(host, port)
        if node_type not in self.type_dict.keys():
            self.type_dict[node_type] = [name]
        else:
            if name not in self.type_dict[node_type]:
                self.type_dict[node_type].append(name)
        # print(self.type_dict, id(self), id(self.type_dict))

    def del_type_node(self, node_type: str, host: str, port: int):
        if isinstance(node_type, NodeType):
            node_type = node_type.value
        name = RpcConnectionManager.gen_node_name(host, port)
        if node_type in self.type_dict.keys():
            if name in self.type_dict[node_type]:
                self.type_dict[node_type].remove(name)

    async def send_message(self, node_name, command_id, data,
                           session_id=None, to=None,
                           msg_type=RpcMsgType.REQUEST):
        """
        向指定连接发送消息
        :param msg_type:
        :param to:
        :param session_id:
        :param data:
        :param node_name: string
        :param command_id: int
        :return:
        """
        if node_name not in self.conns.keys():
            self.warn("can not find connection", node_name=node_name)
            raise ObException(name="ConnectionNotFound", node_name=node_name)

        if ConnectionStatus.ESTABLISHED != self.conns[node_name]["status"]:
            # 连接断开状态
            self.warn("connection is unavailable", node_name=node_name)
            for reconnect_times in range(3):
                await asyncio.sleep(5)
                conn_status = self.conns[node_name]["status"]
                if ConnectionStatus.ESTABLISHED == conn_status:
                    break
                if 2 == reconnect_times:
                    self.error("exceed max reconnect times", node_name=node_name)
                    raise ObException(name="ReconnectMaxTimes",
                                      node_name=node_name)

        # if not isinstance(data, str):
        #     data = ujson.dumps(data)
        # print(data, data.encode("utf8"), type(data)
        result = await self.conns[node_name]["conn"].send_message(
            command_id, data, session_id=session_id, to=to, msg_type=msg_type)
        self.debug("rpc server send end", data=data, session_id=session_id,
                   to=to, result=result, msg_type=msg_type)
        return result
